from resnet import *
from modules import *


class model(nn.Module):
    def __init__(self):
        super(model, self).__init__()
        self.gaze_network = resnet50(pretrained=True)

        # self.gaze_fc = nn.Sequential(
        #     nn.Linear(2048, 2),
        # )
        self.gazeEs = ResGazeEs()

    def forward(self, x):
        feature = self.gaze_network(x)
        # feature = feature.view(feature.size(0), -1)
        # gaze = self.gaze_fc(feature)
        gaze = self.gazeEs(feature)

        return gaze
